package vpc.ext.stacksize.visitors;

import vpc.tir.*;
import vpc.tir.opt.TIRCallShape;
import vpc.tir.opt.DepthFirstOpVisitor;
import vpc.tir.expr.Operator;
import vpc.core.Program;
import vpc.core.types.TypeRef;
import vpc.core.types.Type;
import vpc.core.decl.Method.Temporary;
import vpc.core.decl.Method;
import vpc.core.decl.Method.Temporary;
import vpc.core.ProgramDecl.EntryPoint;
import vpc.core.types.Type;
import vpc.hil.*;
import vpc.hil.DeviceDecl.GetRegister;
import vpc.hil.DeviceDecl.SetRegister;
import vpc.ext.stacksize.*;
import vpc.ext.stacksize.util.Debug;
import vpc.ext.stacksize.util.DataHandler;
import vpc.ext.stacksize.util.UnhandledTypeException;

import java.util.*;

public class GraphVisitor extends
		DepthFirstOpVisitor<TIRExpr, GraphVisitor.Context> {

	private static Debug.Flag debugFlag = Debug.Flag.GRAPHVISITOR;

	public class Context {
		public Node root;

		public Node parent;

		public Context(Node root) {
			this.root = root;
			this.parent = root;
		}
	}

	public final Context context;

	public final List<Node> subRoots;

	private final Program p;

	public GraphVisitor(Program p) {
		super(null);
		Method m = p.programDecl.mainEntry.method;
		Node root = Node.constructMethodNode(m);
		context = new Context(root);
		subRoots = new ArrayList<Node>();
		subRoots.add(root);
		this.p = p;
		
		for (Register r : p.targetDevice.registers) {
			Debug.print("REGISTER: ", debugFlag);
			Debug.print("name = " + r.getName() + " : ", debugFlag);
			Debug.println("token = " + r.index, debugFlag);
		}
		
		for (Interrupt i : p.targetDevice.interrupts) {
			Debug.print("INTERRUPT : ", debugFlag);
			Debug.print("Flagbit = " + i.flagBit + " : ", debugFlag);
			Debug.print("FlagReg = " + i.flagReg + " : ", debugFlag);
			Debug.print("Maskbit = " + i.maskBit + " : ", debugFlag);
			Debug.println("MaskReg = " + i.maskReg, debugFlag);
		}
		
		
	}

	public void setParentByName(String name) {
		this.context.parent = null;

		for (Node n : this.subRoots) {
			if (n.name.equals(name)) {
				this.context.parent = n;
				break;
			}

			Node result = n.find(name, new HashSet<Node>());

			if (result != null) {
				this.context.parent = result;
				break;
			}
		}
	}

	/**
	 * Visit method for TIRExpr that is returned as the result of the getBody()
	 * method.
	 * 
	 * @param body
	 *            TIRExpr representing the body of a method.
	 * @return
	 */
	public TIRExpr visitBody(TIRExpr body) {
		return body.accept(this, this.context);
	}

	public TIRExpr visitOperator(Operator o, Context c) {

		if (o instanceof GetRegister) {
			GetRegister getReg = (GetRegister) o;
			Debug.println("GET_REG: " + getReg.register.getName(), debugFlag);
		}

		else if (o instanceof SetRegister) {
			SetRegister setReg = (SetRegister) o;
			Debug.println("SET_REG: " + setReg.register.getName(), debugFlag);
		}

		return null;
	}

	/**
	 * Visit method for TIROperator type. Visit each sub-operand.
	 */
	public TIRExpr visit(TIROperator o, Context c) {

		this.visitOperator(o.operator, c);
		for (TIRExpr e : o.operands)
			e.accept(this, c);

		return null;
	}

	/**
	 * Visit method for TIRConst.Symbol type. Represents a String, essentially.
	 * {@link vpc.tir.TIRConst.Symbol}
	 */
	public TIRExpr visit(TIRConst.Symbol s, Context c) {

		return null;
	}

	/**
	 * Visit method for TIRConst.Value type. {@link vpc.tir.TIRConst.Value}
	 */
	public TIRExpr visit(TIRConst.Value v, Context c) {

		return null;
	}

	/**
	 * Visit method for TIRLocal.Get type. {@link vpc.tir.TIRLocal.Get}
	 */
	public TIRExpr visit(TIRLocal.Get g, Context c) {

		return null;
	}

	/**
	 * Visit method for TIRLocal.Set type. {@link vpc.tir.TIRLocal.Set}
	 */
	public TIRExpr visit(TIRLocal.Set l, Context c) {
		l.value.accept(this, c);
		return null;
	}

	/**
	 * Visit method for TIRCall type. NOTE: extremely important for call graph
	 * construction.
	 * 
	 * @author Philip Russell
	 * @author Micah Wendell {@link vpc.tir.TIRCall}
	 */
	public TIRExpr visit(TIRCall b, Context c) {

		Debug.Flag localFlag = Debug.Flag.VISIT_TIRCALL;

		// Need to visit the argument TIRExpr(s) as well
		for (TIRExpr arg : b.arguments) {
			arg.accept(this, this.context);
		}

		TIRCallShape.CallShape cs = TIRCallShape.match(p, b);

		List<Method> methodsInvoked = new ArrayList<Method>();

		if (cs instanceof TIRCallShape.Direct) {
			methodsInvoked.add(((TIRCallShape.Direct) cs).method);
		} else if (cs instanceof TIRCallShape.Dynamic) {

			for (Method m : p.closure.methods) {
				TypeRef[] params = m.getArgumentTypes();
				int numArgs = b.arguments.length;

				if (params.length != numArgs)
					continue;

				if (b.getType() != m.getType())
					continue;

				for (int i = 0; i < numArgs; i++) {
					if (params[i].getType() != b.arguments[i].getType())
						continue;
				}

				methodsInvoked.add(m);
			}
		} else { // Virtual

			// FUTURE TODO: add Methods to the methodsInvoked list using CHA

			Method toCall = ((TIRCallShape.Virtual) cs).method;
			TypeRef[] toCallParams = toCall.getArgumentTypes();
			int numArgs = toCallParams.length;

			for (Method m : p.closure.methods) {
				TypeRef[] params = m.getArgumentTypes();

				if (!m.getName().equals(toCall.getName()))
					continue;

				if (params.length != numArgs)
					continue;

				if (b.getType() != m.getType())
					continue;

				for (int i = 0; i < numArgs; i++) {
					if (params[i].getType() != toCallParams[i].getType())
						continue;
				}

				methodsInvoked.add(m);
			}
		}
		for (Method m : methodsInvoked) {

			Node child = null;
			for (Node n : this.subRoots) {

				Debug.println("\tn.id = " + n.name + "; m.getFullName() = "
						+ m.getFullName(), debugFlag);

				if (n.name.equals(m.getFullName())) {

					Debug.println("\tFound matching node in subroots, id = "
							+ n.name, localFlag);

					child = n;
					if (n != context.root && n != c.parent.root) {
						this.subRoots.remove(n);
						n.changeRoot(c.parent.root, new HashSet<String>());
					}
					break;
				}

				child = n.find(m.getFullName(), new HashSet<Node>());

				Debug.println("\tResult of find(): " + child, localFlag);

				if (child != null)
					break;
			}

			if (child == null) {
				child = new Node(m.getFullName(), DataHandler.getStackPush(m),
						c.parent.root);

				Debug.println("\tCreating " + m.getFullName() + " in subgraph",
						localFlag);
			}

			c.parent.addEdge(child);

			Debug.println("\tAdded edge to " + m.getFullName(), localFlag);
		}

		return null;
	}

	/**
	 * Generic visit method for TIRExpr(s) not listed above.
	 * {@link vpc.tir.TIRExpr}
	 */
	public TIRExpr visit(TIRExpr e, Context c) {

		Debug.println("WARNING: visited unknown TIRExpr, CLASS = "
				+ e.getClass().toString(), debugFlag);
		return null;
	}
}
